##########################################################################################

library(data.table)
library(optparse)
library(pheatmap)
library(GSVA)

##########################################################################################

option_list <- list(
    make_option(c("--data_type"), type = "character"),
    make_option(c("--geneset_type"), type = "character"),
    make_option(c("--geneset_file"), type = "character"),
    make_option(c("--rna_file"), type = "character"),
    make_option(c("--cor_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

    data_type <- "exp"
    geneset_type <- "known_motif"
    rna_file <- "~/20231121_singleMuti/results/qc_atac_v3/germ/GeneExpression.MeanByCellType.tsv"
    cor_file <- "~/20231121_singleMuti/results/celltype_plot/mfuzz/cor.motif_atac-rna.tsv"
    out_path <- "~/20231121_singleMuti/results/celltype_plot/mfuzz"
    geneset_file <- "~/20231121_singleMuti/config/Human_reported_TF2.csv"

}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

data_type <- opt$data_type
rna_file <- opt$rna_file
cor_file <- opt$cor_file
out_path <- opt$out_path
geneset_file <- opt$geneset_file
geneset_type <- opt$geneset_type

dir.create(out_path , recursive = T)

###########################################################################################
## 导入数据
sco_exp <- data.frame(fread(rna_file))
dat_cor <- data.frame(fread(cor_file))

##########################################################################################

cell_order <- c("SSC" , "Differenting.Differented.SPG" , "Leptotene" ,
    "Zygotene" , "Patchytene" , "Diplotene" , 
    "Early.stage.of.spermatids" , "Round.ElongateS.tids" , "Sperm"
    )


##########################################################################################

rownames(sco_exp) <- sco_exp$gene
sco_exp <- sco_exp[,-1]

##########################################################################################
## 提取感兴趣的基因集合
if(geneset_type == "all_motif"){
    dat_geneset <- data.frame(fread(geneset_file , header = T))
    tf_name <- dat_geneset$MotifMatrix_matchName
}else if(geneset_type == "known_motif"){
    dat_geneset <- data.frame(fread(geneset_file , header = F))
    tf_name <- dat_geneset$V1
    tf_name <- tf_name[tf_name %in% rownames(sco_exp)]
}

sco_exp <- sco_exp[tf_name,]

##########################################################################################

sco_exp <- sco_exp[which(rowSums(sco_exp)>0),]
sco_exp <- sco_exp[,cell_order]
sco_exp <- t(sco_exp)

for( clust in 2:20){
    p <- pheatmap(sco_exp,
        scale = "column",
        show_rownames = T ,show_colnames = F, 
        #cutree_cols=7,cutree_rows = 7,
        cluster_rows = F , cluster_cols = T , clustering_method = "ward.D2",
        color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
        cutree_cols = clust ,
        cellheight = 10)

    out_file <- paste0( out_path , "/pheatmap_" , data_type , ".cluster." , clust , ".pdf" )
    pdf(out_file)
    print(p)
    dev.off()

    ##########################################################################################
    ## 每类基因进行gsva

    col_cluster <- cutree(p$tree_col,k=clust)
    gsva_path <- data.frame( gene = names(col_cluster) , group_name = col_cluster )
    gsva_path_out <- merge( gsva_path , dat_cor , by.x = "gene" , by.y = "GeneExpressionMatrix_name" , all.x = T )

    out_file <- paste0( out_path , "/pheatmap_" , data_type , ".cluster." , clust , ".gene.tsv" )
    write.table(gsva_path_out, out_file , row.names = F , quote = F , sep = "\t")

    tmp <- list()
    for (i in unique(gsva_path$group_name)){
      i <- list(subset(gsva_path,group_name==i)$gene)
      tmp <- append(tmp,i)
    }

    names(tmp) <- unique(gsva_path$group_name)
    gsva_score <- gsva(as.matrix(t(sco_exp)),tmp,method="ssgsea",verbose=T,parallel.sz=40)

    out_file <- paste0( out_path , "/pheatmap_" , data_type , ".cluster." , clust , ".gsva.tsv" )
    write.table(gsva_score, out_file , row.names = T , quote = F , sep = "\t")

    p <- pheatmap(t(gsva_score),
    scale = "column",
    show_rownames = T ,show_colnames = T, 
    #cutree_cols=7,cutree_rows = 7,
    cluster_rows = F , cluster_cols = T , clustering_method = "ward.D2",
    color = colorRampPalette(c("blue", "white", "firebrick3"))(100),
    cellheight = 10)
    out_file <- paste0( out_path , "/pheatmap_" , data_type , ".cluster." , clust , ".gsva.pdf" )
    pdf(out_file)
    print(p)
    dev.off()

}